#!/usr/bin/env python3
"""
Simple script to plot Value columns from multiple CSV files
"""

import pandas as pd
import matplotlib.pyplot as plt
import os
import glob

def plot_csv_values(data_dir="/home/shaonian/SED/SED/script_tools/data"):
    """
    Plot Value column from all CSV files containing 'conficence' in the name
    
    Args:
        data_dir (str): Directory containing the CSV files
    """
    
    # Find all CSV files with 'conficence' in the name
    pattern = os.path.join(data_dir, "*conficence*.csv")
    csv_files = glob.glob(pattern)
    
    if not csv_files:
        print(f"No CSV files found matching pattern: {pattern}")
        return
    
    print(f"Found {len(csv_files)} CSV files to plot:")
    
    # Create the plot
    plt.figure(figsize=(14, 8))
    
    # Define colors and line styles
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b']
    line_styles = ['-', '--', '-.', ':']
    
    for i, csv_file in enumerate(csv_files):
        try:
            # Read the CSV file
            df = pd.read_csv(csv_file)
            
            # Extract meaningful label from filename
            filename = os.path.basename(csv_file)
            print(f"  Processing: {filename}")
            
            # Create label
            if "teacher" in filename.lower():
                model = "Teacher"
            elif "student" in filename.lower():
                model = "Student"
            else:
                model = "Unknown"
            
            if "in-domain" in filename.lower():
                dataset = "In-domain"
            elif "reftrainweak" in filename.lower():
                dataset = "RefTrainWeak"
            else:
                dataset = "Standard"
            
            label = f"{model} ({dataset})"
            
            # Plot the data
            plt.plot(df['Step'], df['Value'], 
                    label=label,
                    color=colors[i % len(colors)],
                    linestyle=line_styles[i % len(line_styles)],
                    linewidth=2.5,
                    alpha=0.8)
            
            print(f"    - {len(df)} data points, Value range: {df['Value'].min():.3f} - {df['Value'].max():.3f}")
            
        except Exception as e:
            print(f"Error processing {csv_file}: {e}")
    
    # Customize the plot
    plt.xlabel('Training Step', fontsize=14, fontweight='bold')
    plt.ylabel('Confidence Value', fontsize=14, fontweight='bold')
    plt.title('Training Confidence Evolution Comparison\n(Teacher vs Student Models)', 
              fontsize=16, fontweight='bold', pad=20)
    
    # Improve legend
    plt.legend(loc='lower right', fontsize=12, frameon=True, 
              fancybox=True, shadow=True, framealpha=0.9)
    
    # Add grid
    plt.grid(True, alpha=0.3, linestyle='-', linewidth=0.5)
    
    # Set limits
    plt.xlim(left=0)
    plt.ylim(0, 1.0)
    
    # Add some styling
    plt.gca().spines['top'].set_visible(False)
    plt.gca().spines['right'].set_visible(False)
    plt.gca().spines['left'].set_linewidth(1.5)
    plt.gca().spines['bottom'].set_linewidth(1.5)
    
    # Adjust layout
    plt.tight_layout()
    
    # Save the plot
    output_file = os.path.join(data_dir, "confidence_values_plot.png")
    plt.savefig(output_file, dpi=300, bbox_inches='tight', 
                facecolor='white', edgecolor='none')
    
    print(f"\n✅ Plot saved successfully to: {output_file}")
    
    # Display the plot (if running in an environment that supports it)
    plt.show()
    
    return output_file

if __name__ == "__main__":
    # Run the plotting function
    plot_csv_values()
